import tushare as ts
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# 设置Tushare的token
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

def get_stock_data(ts_code, start_date, end_date, freq='D'):
    """
    获取指定股票在指定日期范围内的K线数据
    :param ts_code: 股票代码
    :param start_date: 开始日期，格式：YYYYMMDD
    :param end_date: 结束日期，格式：YYYYMMDD
    :param freq: 数据频率，D-日线，M-月线，W-周线，min-分钟线
    :return: 包含K线数据的DataFrame
    """
    if freq == 'D':
        df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
    elif freq == 'W':
        df = pro.weekly(ts_code=ts_code, start_date=start_date, end_date=end_date)
    elif freq == 'M':
        df = pro.monthly(ts_code=ts_code, start_date=start_date, end_date=end_date)
    elif 'min' in freq:
        df = ts.pro_bar(ts_code=ts_code, start_date=start_date, end_date=end_date, freq=freq)
    else:
        raise ValueError("Unsupported frequency")
    df = df.sort_values(by='trade_date')
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date', inplace=True)
    return df

def generate_signals(df):
    """
    生成交易信号，这里简单使用简单移动平均线交叉策略
    :param df: 包含K线数据的DataFrame
    :return: 包含交易信号的DataFrame
    """
    short_window = 5
    long_window = 20
    df['short_mavg'] = df['close'].rolling(window=short_window, min_periods=1, center=False).mean()
    df['long_mavg'] = df['close'].rolling(window=long_window, min_periods=1, center=False).mean()
    df['signal'] = 0.0
    df['signal'][short_window:] = np.where(df['short_mavg'][short_window:] > df['long_mavg'][short_window:], 1.0, 0.0)
    df['positions'] = df['signal'].diff()
    return df

def backtest(df, initial_capital=100000):
    """
    回测交易策略
    :param df: 包含交易信号的DataFrame
    :param initial_capital: 初始资金
    :return: 包含回测结果的DataFrame
    """
    positions = pd.DataFrame(index=df.index).fillna(0.0)
    positions['stock'] = df['signal']
    portfolio = positions.multiply(df['close'], axis=0)
    pos_diff = positions.diff()
    portfolio['holdings'] = (positions.multiply(df['close'], axis=0)).sum(axis=1)
    portfolio['cash'] = initial_capital - (pos_diff.multiply(df['close'], axis=0)).sum(axis=1).cumsum()
    portfolio['total'] = portfolio['cash'] + portfolio['holdings']
    portfolio['returns'] = portfolio['total'].pct_change()
    return portfolio

def visualize(df, portfolio, title):
    """
    可视化K线、交易信号（买点）、卖点、收益率曲线
    :param df: 包含K线数据的DataFrame
    :param portfolio: 包含回测结果的DataFrame
    :param title: 图表标题
    """
    fig, axes = plt.subplots(2, 1, figsize=(12, 8))
    # 绘制K线和交易信号
    axes[0].plot(df['close'], label='Close Price')
    axes[0].plot(df['short_mavg'], label='Short MA')
    axes[0].plot(df['long_mavg'], label='Long MA')
    axes[0].plot(df[df['positions'] == 1].index, df['short_mavg'][df['positions'] == 1], '^', markersize=10, color='g', label='Buy Signal')
    axes[0].plot(df[df['positions'] == -1].index, df['short_mavg'][df['positions'] == -1], 'v', markersize=10, color='r', label='Sell Signal')
    axes[0].set_title(f'{title} - Price and Signals')
    axes[0].set_xlabel('Date')
    axes[0].set_ylabel('Price')
    axes[0].legend()
    # 绘制收益率曲线
    axes[1].plot(portfolio['total'], label='Portfolio Value')
    axes[1].set_title(f'{title} - Portfolio Value')
    axes[1].set_xlabel('Date')
    axes[1].set_ylabel('Value')
    axes[1].legend()
    plt.tight_layout()
    plt.show()

def multi_stock_backtest(ts_codes, start_date, end_date, freq='D'):
    """
    对多支股票进行回测并比较收益率曲线
    :param ts_codes: 股票代码列表
    :param start_date: 开始日期，格式：YYYYMMDD
    :param end_date: 结束日期，格式：YYYYMMDD
    :param freq: 数据频率，D-日线，M-月线，W-周线，min-分钟线
    """
    plt.figure(figsize=(12, 6))
    for ts_code in ts_codes:
        df = get_stock_data(ts_code, start_date, end_date, freq)
        df = generate_signals(df)
        portfolio = backtest(df)
        plt.plot(portfolio['total'], label=ts_code)
        visualize(df, portfolio, ts_code)
    plt.title('Portfolio Value Comparison')
    plt.xlabel('Date')
    plt.ylabel('Value')
    plt.legend()
    plt.show()

# 示例调用
if __name__ == "__main__":
    ts_codes = ['000001.SZ', '600000.SH']
    start_date = '20200101'
    end_date = '20241231'
    multi_stock_backtest(ts_codes, start_date, end_date, freq='D')